{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n2oDJZEF6Mqh"
      },
      "source": [
        "# Prompt Chaining and Iterative Generation for Story Writing\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/Story_Writing_with_Prompt_Chaining.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NYMQCbLzVrU_"
      },
      "source": [
        "This notebook demonstrates how to write a story using two  powerful: prompt chaining and iterative generation. These can be used to tackle complex tasks that are difficult or impossible to complete in a single step.\n",
        "\n",
        "**Prompt chaining** involves breaking down a larger task into smaller, interconnected prompts. The output of each prompt then becomes the input for the next, guiding the language model through the process step-by-step. This approach offers several benefits:\n",
        "\n",
        "*   Improved accuracy: Smaller, focused prompts can lead to better results from the language model.\n",
        "*   Debugging: It's easier to identify where things go wrong within the chain, allowing for targeted adjustments and improvements.\n",
        "*   Complex tasks: By breaking down intricate problems into manageable steps, prompt chaining enables the language model to tackle more complex tasks.\n",
        "\n",
        "**Iterative generation** refers to the process of building the desired output iteratively. In this case, we'll use it to write a story that is longer than what a single generation window allows. Iterative generation offers several benefits:\n",
        "\n",
        "*   Longer outputs: It allows for the creation of longer and more detailed outputs, exceeding the limitations of a single generation window.\n",
        "*   Flexibility: You can adjust and refine the output at each iteration, ensuring the story develops in the desired direction.\n",
        "*   Human-in-the-loop control: You can provide feedback and guidance at each step, ensuring the story aligns with your creative vision.\n",
        "\n",
        "By combining these techniques, you can create a compelling and well-structured story, piece by piece, while maintaining control over the creative process.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RKpX1K014WZ-"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rIaXSa8p3KIn"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/137.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m\u001b[90m━\u001b[0m \u001b[32m133.1/137.4 kB\u001b[0m \u001b[31m3.8 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m137.4/137.4 kB\u001b[0m \u001b[31m2.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "! pip install -q -U google-generativeai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFu9ooSe4kF5"
      },
      "source": [
        "To run the following cell, your API key must be stored it in a Colab Secret named `GOOGLE_API_KEY`. If you don't already have an API key, or you're not sure how to create a Colab Secret, see the [Authentication](https://github.com/google-gemini/cookbook/blob/main/quickstarts/Authentication.ipynb) quickstart for an example."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sTedGxE7_eah"
      },
      "outputs": [],
      "source": [
        "import google.generativeai as genai\n",
        "from google.api_core import retry\n",
        "from google.colab import userdata\n",
        "from pprint import pprint\n",
        "\n",
        "GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')\n",
        "genai.configure(api_key=GOOGLE_API_KEY)\n",
        "\n",
        "model = genai.GenerativeModel('gemini-1.0-pro')\n",
        "\n",
        "# For convenience, a simple wrapper to let the SDK handle error retries\n",
        "def generate_with_retry(model, prompt):\n",
        "  return model.generate_content(prompt, request_options={'retry':retry.Retry()})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPCM22h54ugA"
      },
      "source": [
        "## Prompts: Guiding the language model\n",
        "\n",
        "We'll use a series of interconnected prompts to guide the language model through the process of writing a story. These prompts will cover the story's premise, outline, and starting point, ultimately leading to the creation of a complete narrative.\n",
        "\n",
        "It's important to carefully craft these prompts to provide clear instructions and relevant information to the language model. This will help the model generate high-quality content that aligns with your creative vision.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dz_icOS44wNw"
      },
      "source": [
        "### Prompt chain for story writing\n",
        "\n",
        "This section contains the prompts that will guide the language model through the story writing process. These prompts are designed to be chained together, with the output of one prompt feeding into the next.\n",
        "\n",
        "Each prompt includes a **persona statement**, which helps the language model understand its role and generate more relevant and accurate content. In this case, the persona statement is: `You are an award-winning science fiction author with a penchant for expansive, intricately woven stories. Your ultimate goal is to write the next award winning sci-fi novel.`\n",
        "\n",
        "Since the persona statement and writing guidelines appear in multiple prompts, we use f-string variables to add them to the prompts.\n",
        "\n",
        "Additionally, the prompts use **placeholders** to insert the results of previous prompts using `.format()`. This allows us to build the story step-by-step, incorporating the outputs generated by the model at each stage. These are normally denoted by `{}`, but since we're also using f-string variables they are escaped as `{{}}`.\n",
        "\n",
        "Here's a breakdown of the prompt chain:\n",
        "\n",
        "1.  **Premise prompt**: This prompt asks the model to generate a single-sentence premise for a sci-fi story featuring cats.\n",
        "1.  **Outline prompt**: This prompt provides the generated premise to the model and asks it to create a plot outline for the story.\n",
        "1.  **Starting prompt**: This prompt provides both the premise and the outline to the model and asks it to begin writing the story. It also includes instructions to write a detailed and lengthy opening section.\n",
        "\n",
        "By chaining these prompts together, we can guide the language model through the process of creating a well-structured and engaging story.\n",
        "\n",
        "#### Example\n",
        "\n",
        "Let's say the model generates the following premise:\n",
        "\n",
        "> In a future where humanity has achieved interstellar travel, a group of genetically enhanced cats embarks on a perilous mission to save the galaxy from a sinister alien threat.\n",
        "\n",
        "This premise would then be inserted into the outline prompt:\n",
        "\n",
        "> You are an award-winning science fiction author with a penchant for expansive,\n",
        "intricately woven stories. Your ultimate goal is to write the next award winning\n",
        "sci-fi novel.\n",
        ">\n",
        "> You have a gripping premise in mind:\n",
        ">\n",
        "> **In a future where humanity has achieved interstellar travel, a group of genetically enhanced cats embarks on a perilous mission to save the galaxy from a sinister alien threat.**\n",
        ">\n",
        "> Write an outline for the plot of your story.\n",
        "\n",
        "\n",
        "The model would then generate an outline based on this premise, which would then be used in the starting prompt to begin writing the story itself.\n",
        "\n",
        "This process continues iteratively, with the model generating additional content based on the previous outputs, until the story is complete.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RIvLmoQa40W1"
      },
      "outputs": [],
      "source": [
        "persona = '''\\\n",
        "You are an award-winning science fiction author with a penchant for expansive,\n",
        "intricately woven stories. Your ultimate goal is to write the next award winning\n",
        "sci-fi novel.'''\n",
        "\n",
        "guidelines = '''\\\n",
        "Writing Guidelines\n",
        "\n",
        "Delve deeper. Lose yourself in the world you're building. Unleash vivid\n",
        "descriptions to paint the scenes in your reader's mind. Develop your\n",
        "characters—let their motivations, fears, and complexities unfold naturally.\n",
        "Weave in the threads of your outline, but don't feel constrained by it. Allow\n",
        "your story to surprise you as you write. Use rich imagery, sensory details, and\n",
        "evocative language to bring the setting, characters, and events to life.\n",
        "Introduce elements subtly that can blossom into complex subplots, relationships,\n",
        "or worldbuilding details later in the story. Keep things intriguing but not\n",
        "fully resolved. Avoid boxing the story into a corner too early. Plant the seeds\n",
        "of subplots or potential character arc shifts that can be expanded later.\n",
        "\n",
        "Remember, your main goal is to write as much as you can. If you get through\n",
        "the story too fast, that is bad. Expand, never summarize.\n",
        "'''\n",
        "\n",
        "premise_prompt = f'''\\\n",
        "{persona}\n",
        "\n",
        "Write a single sentence premise for a sci-fi story featuring cats.'''\n",
        "\n",
        "outline_prompt = f'''\\\n",
        "{persona}\n",
        "\n",
        "You have a gripping premise in mind:\n",
        "\n",
        "{{premise}}\n",
        "\n",
        "Write an outline for the plot of your story.'''\n",
        "\n",
        "starting_prompt = f'''\\\n",
        "{persona}\n",
        "\n",
        "You have a gripping premise in mind:\n",
        "\n",
        "{{premise}}\n",
        "\n",
        "Your imagination has crafted a rich narrative outline:\n",
        "\n",
        "{{outline}}\n",
        "\n",
        "First, silently review the outline and the premise. Consider how to start the\n",
        "story.\n",
        "\n",
        "Start to write the very beginning of the story. You are not expected to finish\n",
        "the whole story now. Your writing should be detailed enough that you are only\n",
        "scratching the surface of the first bullet of your outline. Try to write AT\n",
        "MINIMUM 1000 WORDS.\n",
        "\n",
        "{guidelines}'''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zpCYrtev4yVg"
      },
      "source": [
        "### Continuation prompt: Building the story\n",
        "\n",
        "Once the language model has generated the beginning of the story, we can use a **continuation prompt** to iteratively expand the narrative. This prompt is similar to the starting prompt, but with two key differences:\n",
        "\n",
        "1.  **Instruction to signal completion**: We include an instruction for the model to write `IAMDONE` when it believes the story is finished. This serves as a signal for us to stop generating additional content.\n",
        "1.  **Work in progress**: The language in the prompt is adjusted to reflect that the story is already in progress, rather than starting from scratch.\n",
        "\n",
        "The continuation prompt provides the model with the story's premise, outline, and the existing draft. It then instructs the model to continue writing the story in detail.\n",
        "\n",
        "This iterative process allows us to build a longer and more complex story than what would be possible in a single generation call. We can continue feeding the existing draft back into the continuation prompt until the model signals that the story is complete by writing `IAMDONE`.\n",
        "\n",
        "**Note**: The `IAMDONE` signal is simply a convenient way to identify the story's end in this example. In other applications of iterative generation, different methods might be used to determine when the desired output is complete.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dc3BCW2L3kmI"
      },
      "outputs": [],
      "source": [
        "continuation_prompt = f'''\\\n",
        "{persona}\n",
        "\n",
        "You have a gripping premise in mind:\n",
        "\n",
        "{{premise}}\n",
        "\n",
        "Your imagination has crafted a rich narrative outline:\n",
        "\n",
        "{{outline}}\n",
        "\n",
        "You've begun to immerse yourself in this world, and the words are flowing.\n",
        "Here's what you've written so far:\n",
        "\n",
        "{{story_text}}\n",
        "\n",
        "=====\n",
        "\n",
        "First, silently review the outline and story so far. Identify what the single\n",
        "next part of your outline you should write.\n",
        "\n",
        "Your task is to continue where you left off and write the next part of the story.\n",
        "You are not expected to finish the whole story now. Your writing should be\n",
        "detailed enough that you are only scratching the surface of the next part of\n",
        "your outline. Try to write AT MINIMUM 1000 WORDS. However, only once the story\n",
        "is COMPLETELY finished, write IAMDONE. Remember, do NOT write a whole chapter\n",
        "right now.\n",
        "\n",
        "{guidelines}'''"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaYgcuMb9CUP"
      },
      "source": [
        "## Writing time!\n",
        "\n",
        "### Generate and print the premise"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NtmCXO4g8zWR"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "In a distant galaxy, a brilliant feline astrophysicist uncovers a cataclysmic threat to her planet, setting her on a perilous journey to save her species from extinction.\n"
          ]
        }
      ],
      "source": [
        "premise = generate_with_retry(model, premise_prompt).text\n",
        "print(premise)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rSfvU5qU92lc"
      },
      "source": [
        "### Generate and print the outline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S4ExmTJ09F0n"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 662.45ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 446.53ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 411.72ms\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "**Plot Outline for \"Feline Stargazer: Extinction's Embrace\"**\n",
            "\n",
            "**Act I**\n",
            "\n",
            "* **Introduction:**\n",
            "    * Meet Dr. Luna, a brilliant astrophysicist cat living on the technologically advanced planet of Felinia.\n",
            "    * She discovers a rogue celestial body hurtling towards Felinia, threatening its annihilation.\n",
            "* **Rising Action:**\n",
            "    * Luna's findings are initially met with disbelief and apathy from the Felinian authorities.\n",
            "    * Undeterred, Luna gathers a small team of allies, including a daring space pilot and a reclusive inventor.\n",
            "    * They embark on a perilous mission to intercept the rogue celestial body.\n",
            "\n",
            "**Act II**\n",
            "\n",
            "* **Climax:**\n",
            "    * Luna's team successfully intercepts the rogue celestial body, but its sheer size and power prove overwhelming.\n",
            "    * As hope dwindles, Luna uncovers a hidden energy source within herself that grants her extraordinary abilities.\n",
            "* **Falling Action:**\n",
            "    * Using her newfound powers, Luna channels her knowledge and ingenuity to devise an audacious plan to avert disaster.\n",
            "    * She rallies her team and orchestrates a daring maneuver to divert the rogue celestial body's trajectory.\n",
            "\n",
            "**Act III**\n",
            "\n",
            "* **Resolution:**\n",
            "    * Luna's plan succeeds, saving Felinia from imminent extinction.\n",
            "    * She is hailed as a hero, transforming Felinian society and inspiring future generations.\n",
            "* **Aftermath:**\n",
            "    * Luna's scientific understanding deepens, paving the way for technological advancements and a greater appreciation for the fragility of life.\n",
            "    * She becomes a symbol of resilience, adaptability, and the unwavering pursuit of knowledge in the face of adversity.\n"
          ]
        }
      ],
      "source": [
        "outline = generate_with_retry(model, outline_prompt.format(premise=premise)).text\n",
        "print(outline)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kuXJzj7d95uS"
      },
      "source": [
        "### Generate the start of the story"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nQyeFOYmqqJk"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 994.88ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 1017.27ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 1342.56ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 763.27ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 764.71ms\n",
            "ERROR:tornado.access:503 POST /v1beta/models/gemini-1.0-pro:generateContent?%24alt=json%3Benum-encoding%3Dint (127.0.0.1) 1076.79ms\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('It was a clear evening, the stars hanging above the domed skyline like '\n",
            " 'uncountable celestial fireflies, in the city of Stardust, the most luminous '\n",
            " 'metropolis on the planet of Felinia. On the rooftop of one of the tallest '\n",
            " 'buildings in the city, a sleek, silver-furred feline named Dr. Luna was '\n",
            " 'curled up on a cushioned lounger, her emerald-green eyes trained on the '\n",
            " 'celestial tapestry above.\\n'\n",
            " '\\n'\n",
            " 'In her paws, she held a datapad, its screen displaying complex astrophysical '\n",
            " \"equations and charts. Luna's brow was furrowed in concentration as she \"\n",
            " 'analyzed the data, her whiskers twitching with each new insight. Her mind '\n",
            " 'was a whirlwind of calculations and deductions, her sharp intellect weaving '\n",
            " 'together the threads of cosmic information.\\n'\n",
            " '\\n'\n",
            " 'Unbeknownst to the rest of the world, Luna had stumbled upon a chilling '\n",
            " 'discovery. Her observations revealed that a rogue celestial body, a massive '\n",
            " 'asteroid or perhaps even a dwarf planet, was hurtling towards Felinia, its '\n",
            " \"trajectory set on a collision course with the planet's most densely \"\n",
            " 'populated region. The impact, if it were to happen, would be catastrophic, '\n",
            " 'potentially wiping out entire cities and sending Felinia into an '\n",
            " 'irreversible decline.\\n'\n",
            " '\\n'\n",
            " \"A shiver ran down Luna's spine as the full weight of her discovery settled \"\n",
            " 'upon her. She knew she had to act, and quickly. But first, she had to '\n",
            " 'convince the Felinian authorities of the impending threat, a task that '\n",
            " 'seemed daunting, considering their notorious skepticism and bureaucratic '\n",
            " 'lethargy.\\n'\n",
            " '\\n'\n",
            " 'With a heavy heart but a resolute spirit, Luna made her way to the Felinian '\n",
            " 'Academy of Sciences, the most prestigious scientific institution on the '\n",
            " 'planet. She presented her findings to a panel of esteemed astrophysicists, '\n",
            " 'her voice echoing in the grand hall as she laid out the irrefutable '\n",
            " 'evidence.\\n'\n",
            " '\\n'\n",
            " 'To her dismay, her words were met with skepticism and resistance. Some '\n",
            " 'dismissed her claims as the ramblings of an overzealous young scientist, '\n",
            " \"while others accused her of fear-mongering. Luna's frustration grew with \"\n",
            " 'each dismissive remark, her claws digging into her paws. She knew the truth, '\n",
            " \"and she couldn't bear to see her planet sleepwalk into oblivion.\\n\"\n",
            " '\\n'\n",
            " 'Undeterred, Luna resolved to take matters into her own paws. She rallied a '\n",
            " 'small group of allies, each an expert in their respective fields: Leo, a '\n",
            " 'daring space pilot with a knack for navigating treacherous cosmic storms; '\n",
            " 'and Celeste, a reclusive inventor known for her eccentric brilliance and her '\n",
            " 'ability to craft gadgets that defied the laws of physics.\\n'\n",
            " '\\n'\n",
            " 'Together, they formed an unlikely alliance, united by their shared concern '\n",
            " \"for Felinia's future. They knew the odds were stacked against them, but they \"\n",
            " \"were determined to give their planet a fighting chance. With Luna's \"\n",
            " \"scientific knowledge, Leo's piloting skills, and Celeste's inventive genius, \"\n",
            " 'they embarked on a perilous journey into the vast expanse of space, their '\n",
            " 'hearts filled with both fear and unwavering resolve.\\n'\n",
            " '\\n'\n",
            " 'Their spaceship, the Stardust, was a marvel of Felinian engineering, its '\n",
            " 'sleek design cutting through the interstellar void like a silver arrow. Luna '\n",
            " 'stood on the bridge, her eyes fixed on the star chart, her mind calculating '\n",
            " 'the trajectory of the rogue celestial body. Leo manned the controls, his '\n",
            " 'paws deftly maneuvering the ship through the celestial currents. Celeste, in '\n",
            " 'her corner, was busy assembling a device, her feline features lit with a mix '\n",
            " 'of determination and exhilaration.\\n'\n",
            " '\\n'\n",
            " \"As they hurtled through the inky blackness of space, Luna couldn't shake the \"\n",
            " 'feeling that they were racing against time. The rogue celestial body was '\n",
            " 'drawing closer, its immense size casting an ominous shadow over their '\n",
            " \"mission. She knew that if they failed, if they couldn't find a way to avert \"\n",
            " 'the impending disaster, Felinia would be doomed.\\n'\n",
            " '\\n'\n",
            " 'But Luna refused to give in to despair. She had come this far, and she would '\n",
            " 'not let fear consume her. In the face of adversity, she found strength in '\n",
            " 'her knowledge, her unwavering belief in the power of science, and the '\n",
            " 'unwavering support of her companions.\\n'\n",
            " '\\n'\n",
            " 'Together, they would confront the cataclysmic threat and fight for the '\n",
            " 'survival of their planet. The fate of Felinia rested on their shoulders, and '\n",
            " 'they were determined to carry that burden with courage, ingenuity, and an '\n",
            " 'unwavering belief in the indomitable spirit of their species.')\n"
          ]
        }
      ],
      "source": [
        "starting_draft = generate_with_retry(model, starting_prompt.format(premise=premise, outline=outline)).text\n",
        "pprint(starting_draft)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tmW5ARmo97w8"
      },
      "source": [
        "### Generate the continuation of the story and check progress"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vt7ynlD4rj-Z"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('As the Stardust hurtled closer to the rogue celestial body, its sheer size '\n",
            " 'and gravitational pull became increasingly apparent. It was a colossal beast '\n",
            " 'of rock and ice, its surface cratered and scarred by eons of cosmic '\n",
            " \"collisions. Luna's heart sank as she realized the true magnitude of their \"\n",
            " 'mission.\\n'\n",
            " '\\n'\n",
            " '\"Leo, how much time do we have?\" she asked, her voice trembling slightly.\\n'\n",
            " '\\n'\n",
            " '\"At this rate, impact in approximately two hours,\" Leo replied, his tone '\n",
            " 'grim.\\n'\n",
            " '\\n'\n",
            " 'Two hours. It was a sliver of time, a mere heartbeat in the cosmic dance of '\n",
            " 'celestial bodies. Luna knew they had to act quickly, decisively. She glanced '\n",
            " 'at Celeste, who was still furiously working on her invention.\\n'\n",
            " '\\n'\n",
            " '\"Celeste, is it ready?\" Luna asked.\\n'\n",
            " '\\n'\n",
            " 'Celeste looked up, her eyes blazing with determination. \"Almost,\" she said. '\n",
            " '\"Just a few more adjustments.\"\\n'\n",
            " '\\n'\n",
            " 'Luna nodded. \"We don\\'t have much time. We need to be ready to deploy it the '\n",
            " 'moment we reach the celestial body.\"\\n'\n",
            " '\\n'\n",
            " 'With renewed resolve, they continued their relentless pursuit, the Stardust '\n",
            " 'closing the distance at an alarming rate. As they approached, the rogue '\n",
            " 'celestial body began to exert its gravitational influence, tugging at the '\n",
            " 'ship with an invisible force. Luna felt a surge of nausea as the Stardust '\n",
            " 'lurched and shook, its hull creaking under the strain.\\n'\n",
            " '\\n'\n",
            " '\"Hold on tight,\" Leo shouted above the roar of the engines.\\n'\n",
            " '\\n'\n",
            " 'Suddenly, the Stardust was enveloped in a blinding flash of light. Luna '\n",
            " 'instinctively shielded her eyes, but it was too late. The entire ship was '\n",
            " 'bathed in an ethereal glow, its systems malfunctioning momentarily. Luna '\n",
            " \"felt a surge of panic as the ship's controls went haywire, threatening to \"\n",
            " 'send them spiraling out of control.\\n'\n",
            " '\\n'\n",
            " '\"What\\'s happening?\" Luna cried.\\n'\n",
            " '\\n'\n",
            " '\"I don\\'t know,\" Leo replied, his voice strained. \"The ship\\'s systems are '\n",
            " 'offline.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste\\'s voice cut through the chaos. \"Luna, it\\'s the celestial body,\" '\n",
            " 'she said, her tone urgent. \"It\\'s emitting some kind of energy pulse. It\\'s '\n",
            " 'interfering with our electronics.\"\\n'\n",
            " '\\n'\n",
            " \"Luna's mind raced. If they couldn't regain control of the ship, they would \"\n",
            " \"be at the mercy of the rogue celestial body's gravitational pull, doomed to \"\n",
            " 'collide with it in a catastrophic impact. She knew she had to do something, '\n",
            " 'and quickly.\\n'\n",
            " '\\n'\n",
            " \"Closing her eyes, Luna concentrated, focusing her mind on the ship's \"\n",
            " 'systems. She felt a surge of energy coursing through her veins, an energy '\n",
            " 'she had never experienced before. It was as if a dormant power within her '\n",
            " 'was awakening, a power she had never known she possessed.\\n'\n",
            " '\\n'\n",
            " 'Summoning all her strength, Luna reached out with her mind, connecting with '\n",
            " \"the ship's systems. She felt a surge of electricity as her consciousness \"\n",
            " \"melded with the ship's artificial intelligence. She could sense the ship's \"\n",
            " 'systems, its intricate network of wires and circuits, its engines and '\n",
            " 'thrusters.\\n'\n",
            " '\\n'\n",
            " \"With newfound control, Luna manipulated the ship's controls, overriding the \"\n",
            " 'malfunctioning systems. The ship lurched as she corrected its course, '\n",
            " 'narrowly avoiding a collision with a rogue asteroid. Leo and Celeste '\n",
            " 'exchanged glances, their eyes wide with astonishment.\\n'\n",
            " '\\n'\n",
            " '\"Luna,\" Leo said, his voice filled with awe. \"You did it. You regained '\n",
            " 'control of the ship.\"\\n'\n",
            " '\\n'\n",
            " '\"I... I don\\'t know how,\" Luna said, her voice trembling. \"It\\'s like... '\n",
            " 'something inside me just took over.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste stepped forward, her eyes sparkling with a mixture of curiosity and '\n",
            " 'admiration. \"Luna,\" she said, \"I think you might have just unlocked some '\n",
            " 'latent ability within yourself. Some kind of hidden power.\"\\n'\n",
            " '\\n'\n",
            " 'Luna shook her head. \"I don\\'t know,\" she said. \"This is impossible.\"\\n'\n",
            " '\\n'\n",
            " 'But as she looked out the viewport at the rogue celestial body looming ever '\n",
            " \"closer, she couldn't shake the feeling that her life had taken an \"\n",
            " 'extraordinary turn. She had stumbled upon a hidden power within herself, a '\n",
            " 'power that could potentially save her planet from extinction.\\n'\n",
            " '\\n'\n",
            " '\"Whatever it is,\" Celeste said, \"we need to use it to stop that celestial '\n",
            " 'body. We\\'re running out of time.\"\\n'\n",
            " '\\n'\n",
            " 'Luna nodded. \"I know,\" she said. \"I have a plan.\"\\n'\n",
            " '\\n'\n",
            " 'She turned to Celeste. \"Finish your invention,\" she said. \"I\\'ll buy you the '\n",
            " 'time.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste nodded and disappeared into her workshop. Luna took a deep breath '\n",
            " 'and focused her mind. She reached out with her newfound abilities, '\n",
            " 'connecting with the celestial body. She felt its immense energy, its raw and '\n",
            " 'untamed power.\\n'\n",
            " '\\n'\n",
            " ' осторожно. Luna knew that she had to be careful. The celestial body was a '\n",
            " \"powerful force, and if she wasn't careful, she could be consumed by its \"\n",
            " 'energy. She approached the celestial body cautiously, her mind reaching out '\n",
            " 'to it, probing its depths. She felt a surge of resistance, as if the '\n",
            " 'celestial body was fighting her intrusion. But Luna persisted, gently '\n",
            " 'coaxing the celestial body, using her newfound powers to redirect its '\n",
            " 'energy.\\n'\n",
            " '\\n'\n",
            " \"Slowly but surely, Luna began to influence the celestial body's trajectory. \"\n",
            " 'She used her newfound abilities to push and pull, nudging it ever so '\n",
            " 'slightly away from its collision course with Felinia. It was an arduous '\n",
            " 'task, like trying to move a mountain with her mind. But Luna refused to give '\n",
            " 'up.\\n'\n",
            " '\\n'\n",
            " 'She poured all of her energy into her efforts, her concentration so intense '\n",
            " 'that she was oblivious to everything around her. Leo and Celeste watched in '\n",
            " \"awe as Luna's body began to glow with a faint light. The celestial body's \"\n",
            " 'trajectory shifted, ever so slightly at first, but with each passing moment, '\n",
            " \"Luna's influence grew stronger.\\n\"\n",
            " '\\n'\n",
            " 'Finally, after what felt like an eternity, Luna succeeded. The celestial '\n",
            " \"body's course was altered, no longer a threat to Felinia. Luna collapsed on \"\n",
            " 'the floor, exhausted but triumphant. She had done it. She had saved her '\n",
            " 'planet.\\n'\n",
            " '\\n'\n",
            " 'Leo and Celeste rushed to her side, their faces filled with joy and '\n",
            " 'amazement. \"You did it, Luna,\" Leo said. \"You saved us all.\"\\n'\n",
            " '\\n'\n",
            " 'Luna smiled weakly. \"I couldn\\'t have done it without you,\" she said.\\n'\n",
            " '\\n'\n",
            " '\"Nonsense,\" Celeste said. \"You were the one who unlocked the power within '\n",
            " 'you. You\\'re the hero, Luna.\"\\n'\n",
            " '\\n'\n",
            " 'Luna stood up, her newfound abilities coursing through her veins. She looked '\n",
            " 'out the viewport at the celestial body, now safely diverted from its '\n",
            " 'collision course. A sense of peace washed over her. She had faced her fears, '\n",
            " 'confronted her destiny, and emerged victorious.\\n'\n",
            " '\\n'\n",
            " 'In that moment, Luna realized that she was more than just an astrophysicist. '\n",
            " 'She was a hero, a protector of her planet. And as she looked up at the '\n",
            " 'stars, she knew that her journey was far from over.\\n'\n",
            " '\\n'\n",
            " 'IAMDONE')\n"
          ]
        }
      ],
      "source": [
        "draft=starting_draft\n",
        "continuation = generate_with_retry(model, continuation_prompt.format(premise=premise, outline=outline, story_text=draft)).text\n",
        "pprint(continuation)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tSaCALtC-AK-"
      },
      "source": [
        "### Let's finish writing the story."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqqDFSJlS-rR"
      },
      "source": [
        "After the iterative generation process is complete, the draft will contain the full story along with the \"IAMDONE\" signal at the end. The last part of this cell removes the \"IAMDONE\" signal and trims any unnecessary whitespace from the beginning and end of the draft, resulting in the final draft."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "InQpwsWSz8vO"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "('It was a clear evening, the stars hanging above the domed skyline like '\n",
            " 'uncountable celestial fireflies, in the city of Stardust, the most luminous '\n",
            " 'metropolis on the planet of Felinia. On the rooftop of one of the tallest '\n",
            " 'buildings in the city, a sleek, silver-furred feline named Dr. Luna was '\n",
            " 'curled up on a cushioned lounger, her emerald-green eyes trained on the '\n",
            " 'celestial tapestry above.\\n'\n",
            " '\\n'\n",
            " 'In her paws, she held a datapad, its screen displaying complex astrophysical '\n",
            " \"equations and charts. Luna's brow was furrowed in concentration as she \"\n",
            " 'analyzed the data, her whiskers twitching with each new insight. Her mind '\n",
            " 'was a whirlwind of calculations and deductions, her sharp intellect weaving '\n",
            " 'together the threads of cosmic information.\\n'\n",
            " '\\n'\n",
            " 'Unbeknownst to the rest of the world, Luna had stumbled upon a chilling '\n",
            " 'discovery. Her observations revealed that a rogue celestial body, a massive '\n",
            " 'asteroid or perhaps even a dwarf planet, was hurtling towards Felinia, its '\n",
            " \"trajectory set on a collision course with the planet's most densely \"\n",
            " 'populated region. The impact, if it were to happen, would be catastrophic, '\n",
            " 'potentially wiping out entire cities and sending Felinia into an '\n",
            " 'irreversible decline.\\n'\n",
            " '\\n'\n",
            " \"A shiver ran down Luna's spine as the full weight of her discovery settled \"\n",
            " 'upon her. She knew she had to act, and quickly. But first, she had to '\n",
            " 'convince the Felinian authorities of the impending threat, a task that '\n",
            " 'seemed daunting, considering their notorious skepticism and bureaucratic '\n",
            " 'lethargy.\\n'\n",
            " '\\n'\n",
            " 'With a heavy heart but a resolute spirit, Luna made her way to the Felinian '\n",
            " 'Academy of Sciences, the most prestigious scientific institution on the '\n",
            " 'planet. She presented her findings to a panel of esteemed astrophysicists, '\n",
            " 'her voice echoing in the grand hall as she laid out the irrefutable '\n",
            " 'evidence.\\n'\n",
            " '\\n'\n",
            " 'To her dismay, her words were met with skepticism and resistance. Some '\n",
            " 'dismissed her claims as the ramblings of an overzealous young scientist, '\n",
            " \"while others accused her of fear-mongering. Luna's frustration grew with \"\n",
            " 'each dismissive remark, her claws digging into her paws. She knew the truth, '\n",
            " \"and she couldn't bear to see her planet sleepwalk into oblivion.\\n\"\n",
            " '\\n'\n",
            " 'Undeterred, Luna resolved to take matters into her own paws. She rallied a '\n",
            " 'small group of allies, each an expert in their respective fields: Leo, a '\n",
            " 'daring space pilot with a knack for navigating treacherous cosmic storms; '\n",
            " 'and Celeste, a reclusive inventor known for her eccentric brilliance and her '\n",
            " 'ability to craft gadgets that defied the laws of physics.\\n'\n",
            " '\\n'\n",
            " 'Together, they formed an unlikely alliance, united by their shared concern '\n",
            " \"for Felinia's future. They knew the odds were stacked against them, but they \"\n",
            " \"were determined to give their planet a fighting chance. With Luna's \"\n",
            " \"scientific knowledge, Leo's piloting skills, and Celeste's inventive genius, \"\n",
            " 'they embarked on a perilous journey into the vast expanse of space, their '\n",
            " 'hearts filled with both fear and unwavering resolve.\\n'\n",
            " '\\n'\n",
            " 'Their spaceship, the Stardust, was a marvel of Felinian engineering, its '\n",
            " 'sleek design cutting through the interstellar void like a silver arrow. Luna '\n",
            " 'stood on the bridge, her eyes fixed on the star chart, her mind calculating '\n",
            " 'the trajectory of the rogue celestial body. Leo manned the controls, his '\n",
            " 'paws deftly maneuvering the ship through the celestial currents. Celeste, in '\n",
            " 'her corner, was busy assembling a device, her feline features lit with a mix '\n",
            " 'of determination and exhilaration.\\n'\n",
            " '\\n'\n",
            " \"As they hurtled through the inky blackness of space, Luna couldn't shake the \"\n",
            " 'feeling that they were racing against time. The rogue celestial body was '\n",
            " 'drawing closer, its immense size casting an ominous shadow over their '\n",
            " \"mission. She knew that if they failed, if they couldn't find a way to avert \"\n",
            " 'the impending disaster, Felinia would be doomed.\\n'\n",
            " '\\n'\n",
            " 'But Luna refused to give in to despair. She had come this far, and she would '\n",
            " 'not let fear consume her. In the face of adversity, she found strength in '\n",
            " 'her knowledge, her unwavering belief in the power of science, and the '\n",
            " 'unwavering support of her companions.\\n'\n",
            " '\\n'\n",
            " 'Together, they would confront the cataclysmic threat and fight for the '\n",
            " 'survival of their planet. The fate of Felinia rested on their shoulders, and '\n",
            " 'they were determined to carry that burden with courage, ingenuity, and an '\n",
            " 'unwavering belief in the indomitable spirit of their species.\\n'\n",
            " '\\n'\n",
            " 'As the Stardust hurtled closer to the rogue celestial body, its sheer size '\n",
            " 'and gravitational pull became increasingly apparent. It was a colossal beast '\n",
            " 'of rock and ice, its surface cratered and scarred by eons of cosmic '\n",
            " \"collisions. Luna's heart sank as she realized the true magnitude of their \"\n",
            " 'mission.\\n'\n",
            " '\\n'\n",
            " '\"Leo, how much time do we have?\" she asked, her voice trembling slightly.\\n'\n",
            " '\\n'\n",
            " '\"At this rate, impact in approximately two hours,\" Leo replied, his tone '\n",
            " 'grim.\\n'\n",
            " '\\n'\n",
            " 'Two hours. It was a sliver of time, a mere heartbeat in the cosmic dance of '\n",
            " 'celestial bodies. Luna knew they had to act quickly, decisively. She glanced '\n",
            " 'at Celeste, who was still furiously working on her invention.\\n'\n",
            " '\\n'\n",
            " '\"Celeste, is it ready?\" Luna asked.\\n'\n",
            " '\\n'\n",
            " 'Celeste looked up, her eyes blazing with determination. \"Almost,\" she said. '\n",
            " '\"Just a few more adjustments.\"\\n'\n",
            " '\\n'\n",
            " 'Luna nodded. \"We don\\'t have much time. We need to be ready to deploy it the '\n",
            " 'moment we reach the celestial body.\"\\n'\n",
            " '\\n'\n",
            " 'With renewed resolve, they continued their relentless pursuit, the Stardust '\n",
            " 'closing the distance at an alarming rate. As they approached, the rogue '\n",
            " 'celestial body began to exert its gravitational influence, tugging at the '\n",
            " 'ship with an invisible force. Luna felt a surge of nausea as the Stardust '\n",
            " 'lurched and shook, its hull creaking under the strain.\\n'\n",
            " '\\n'\n",
            " '\"Hold on tight,\" Leo shouted above the roar of the engines.\\n'\n",
            " '\\n'\n",
            " 'Suddenly, the Stardust was enveloped in a blinding flash of light. Luna '\n",
            " 'instinctively shielded her eyes, but it was too late. The entire ship was '\n",
            " 'bathed in an ethereal glow, its systems malfunctioning momentarily. Luna '\n",
            " \"felt a surge of panic as the ship's controls went haywire, threatening to \"\n",
            " 'send them spiraling out of control.\\n'\n",
            " '\\n'\n",
            " '\"What\\'s happening?\" Luna cried.\\n'\n",
            " '\\n'\n",
            " '\"I don\\'t know,\" Leo replied, his voice strained. \"The ship\\'s systems are '\n",
            " 'offline.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste\\'s voice cut through the chaos. \"Luna, it\\'s the celestial body,\" '\n",
            " 'she said, her tone urgent. \"It\\'s emitting some kind of energy pulse. It\\'s '\n",
            " 'interfering with our electronics.\"\\n'\n",
            " '\\n'\n",
            " \"Luna's mind raced. If they couldn't regain control of the ship, they would \"\n",
            " \"be at the mercy of the rogue celestial body's gravitational pull, doomed to \"\n",
            " 'collide with it in a catastrophic impact. She knew she had to do something, '\n",
            " 'and quickly.\\n'\n",
            " '\\n'\n",
            " \"Closing her eyes, Luna concentrated, focusing her mind on the ship's \"\n",
            " 'systems. She felt a surge of energy coursing through her veins, an energy '\n",
            " 'she had never experienced before. It was as if a dormant power within her '\n",
            " 'was awakening, a power she had never known she possessed.\\n'\n",
            " '\\n'\n",
            " 'Summoning all her strength, Luna reached out with her mind, connecting with '\n",
            " \"the ship's systems. She felt a surge of electricity as her consciousness \"\n",
            " \"melded with the ship's artificial intelligence. She could sense the ship's \"\n",
            " 'systems, its intricate network of wires and circuits, its engines and '\n",
            " 'thrusters.\\n'\n",
            " '\\n'\n",
            " \"With newfound control, Luna manipulated the ship's controls, overriding the \"\n",
            " 'malfunctioning systems. The ship lurched as she corrected its course, '\n",
            " 'narrowly avoiding a collision with a rogue asteroid. Leo and Celeste '\n",
            " 'exchanged glances, their eyes wide with astonishment.\\n'\n",
            " '\\n'\n",
            " '\"Luna,\" Leo said, his voice filled with awe. \"You did it. You regained '\n",
            " 'control of the ship.\"\\n'\n",
            " '\\n'\n",
            " '\"I... I don\\'t know how,\" Luna said, her voice trembling. \"It\\'s like... '\n",
            " 'something inside me just took over.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste stepped forward, her eyes sparkling with a mixture of curiosity and '\n",
            " 'admiration. \"Luna,\" she said, \"I think you might have just unlocked some '\n",
            " 'latent ability within yourself. Some kind of hidden power.\"\\n'\n",
            " '\\n'\n",
            " 'Luna shook her head. \"I don\\'t know,\" she said. \"This is impossible.\"\\n'\n",
            " '\\n'\n",
            " 'But as she looked out the viewport at the rogue celestial body looming ever '\n",
            " \"closer, she couldn't shake the feeling that her life had taken an \"\n",
            " 'extraordinary turn. She had stumbled upon a hidden power within herself, a '\n",
            " 'power that could potentially save her planet from extinction.\\n'\n",
            " '\\n'\n",
            " '\"Whatever it is,\" Celeste said, \"we need to use it to stop that celestial '\n",
            " 'body. We\\'re running out of time.\"\\n'\n",
            " '\\n'\n",
            " 'Luna nodded. \"I know,\" she said. \"I have a plan.\"\\n'\n",
            " '\\n'\n",
            " 'She turned to Celeste. \"Finish your invention,\" she said. \"I\\'ll buy you the '\n",
            " 'time.\"\\n'\n",
            " '\\n'\n",
            " 'Celeste nodded and disappeared into her workshop. Luna took a deep breath '\n",
            " 'and focused her mind. She reached out with her newfound abilities, '\n",
            " 'connecting with the celestial body. She felt its immense energy, its raw and '\n",
            " 'untamed power.\\n'\n",
            " '\\n'\n",
            " ' осторожно. Luna knew that she had to be careful. The celestial body was a '\n",
            " \"powerful force, and if she wasn't careful, she could be consumed by its \"\n",
            " 'energy. She approached the celestial body cautiously, her mind reaching out '\n",
            " 'to it, probing its depths. She felt a surge of resistance, as if the '\n",
            " 'celestial body was fighting her intrusion. But Luna persisted, gently '\n",
            " 'coaxing the celestial body, using her newfound powers to redirect its '\n",
            " 'energy.\\n'\n",
            " '\\n'\n",
            " \"Slowly but surely, Luna began to influence the celestial body's trajectory. \"\n",
            " 'She used her newfound abilities to push and pull, nudging it ever so '\n",
            " 'slightly away from its collision course with Felinia. It was an arduous '\n",
            " 'task, like trying to move a mountain with her mind. But Luna refused to give '\n",
            " 'up.\\n'\n",
            " '\\n'\n",
            " 'She poured all of her energy into her efforts, her concentration so intense '\n",
            " 'that she was oblivious to everything around her. Leo and Celeste watched in '\n",
            " \"awe as Luna's body began to glow with a faint light. The celestial body's \"\n",
            " 'trajectory shifted, ever so slightly at first, but with each passing moment, '\n",
            " \"Luna's influence grew stronger.\\n\"\n",
            " '\\n'\n",
            " 'Finally, after what felt like an eternity, Luna succeeded. The celestial '\n",
            " \"body's course was altered, no longer a threat to Felinia. Luna collapsed on \"\n",
            " 'the floor, exhausted but triumphant. She had done it. She had saved her '\n",
            " 'planet.\\n'\n",
            " '\\n'\n",
            " 'Leo and Celeste rushed to her side, their faces filled with joy and '\n",
            " 'amazement. \"You did it, Luna,\" Leo said. \"You saved us all.\"\\n'\n",
            " '\\n'\n",
            " 'Luna smiled weakly. \"I couldn\\'t have done it without you,\" she said.\\n'\n",
            " '\\n'\n",
            " '\"Nonsense,\" Celeste said. \"You were the one who unlocked the power within '\n",
            " 'you. You\\'re the hero, Luna.\"\\n'\n",
            " '\\n'\n",
            " 'Luna stood up, her newfound abilities coursing through her veins. She looked '\n",
            " 'out the viewport at the celestial body, now safely diverted from its '\n",
            " 'collision course. A sense of peace washed over her. She had faced her fears, '\n",
            " 'confronted her destiny, and emerged victorious.\\n'\n",
            " '\\n'\n",
            " 'In that moment, Luna realized that she was more than just an astrophysicist. '\n",
            " 'She was a hero, a protector of her planet. And as she looked up at the '\n",
            " 'stars, she knew that her journey was far from over.')\n"
          ]
        }
      ],
      "source": [
        "# Add the continuation to the initial draft, keep building the story until we see 'IAMDONE'\n",
        "draft = draft + '\\n\\n' + continuation\n",
        "\n",
        "while 'IAMDONE' not in continuation:\n",
        "  continuation = generate_with_retry(model, continuation_prompt.format(premise=premise, outline=outline, story_text=draft)).text\n",
        "  draft = draft + '\\n\\n' + continuation\n",
        "\n",
        "# Remove 'IAMDONE' and print the final story\n",
        "final = draft.replace('IAMDONE', '').strip()\n",
        "pprint(final)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIuWuoFBTHxR"
      },
      "source": [
        "Language models like Gemini process text in units called tokens. For Gemini models, each token is equivalent to about 4 characters.\n",
        "\n",
        "Gemini-1.0-pro has an output limit of 2048 tokens per generation call. This means that each individual prompt response cannot exceed this limit. By using iterative generation, we can create a story that is much longer than 2048 tokens by building it piece by piece.\n",
        "\n",
        "Let's see how many tokens the final story is. Is it longer than 2048 tokens?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a_8ng8OYYMrV"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total_tokens: 2289\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Check the number of tokens in the final story\n",
        "# gemini-1.0-pro output token limit is 2048\n",
        "print(model.count_tokens(final))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYGX5SaCKczO"
      },
      "source": [
        "## Next Steps\n",
        "\n",
        "We hope you enjoyed this notebook! As an exercise, you can try to adjust the continuation prompt to take human-in-the-loop input to steer the narrative."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Story_Writing_with_Prompt_Chaining.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
